from flask import Flask, request, Response  # 1. 导入Response用于构建流式响应
from openai import OpenAI
from Config import config  # 确保Config.py中存在apikey配置
app = Flask(__name__)

# 初始化Kimi客户端（与原代码一致）
client = OpenAI(
    api_key=config['apikey'],
    base_url="https://api.moonshot.cn/v1",
)

# 全局对话历史（生产环境需替换为数据库，避免内存泄漏）
conversation_history = [
    {"role": "system", "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，负责友好、准确地解答用户问题。"}
]

def _build_messages_with_retention(full_history, retain_rounds):
    """历史消息处理函数（原逻辑不变）"""
    if retain_rounds is None:
        return full_history
    system_message = None
    messages_after_system = full_history
    # 提取系统消息（若存在）
    if full_history and full_history[0].get("role") == "system":
        system_message = full_history[0]
        messages_after_system = full_history[1:]
    # 无对话历史时返回系统消息
    if not messages_after_system:
        return [system_message] if system_message else []
    # 拆分当前用户消息与历史消息
    current_user_message = messages_after_system[-1:]
    previous_messages = messages_after_system[:-1]
    # 计算需保留的历史轮次（1轮=用户+助手，故×2）
    keep_count = max(0, int(retain_rounds)) * 2
    kept_previous = previous_messages[-keep_count:] if keep_count > 0 else []
    # 重组消息列表
    messages = []
    if system_message:
        messages.append(system_message)
    messages.extend(kept_previous)
    messages.extend(current_user_message)
    return messages

@app.route('/chat', methods=['POST'])
def chat_endpoint():
    """流式聊天API接口（核心修改部分）"""
    global conversation_history
    
    # 1. 获取请求参数（与原代码一致）
    data = request.json
    query = data.get('message')
    retain_rounds = data.get('retain_rounds', 5)  # 默认保留5轮对话
    if not query:
        return {"error": "Missing 'message' parameter"}, 400
    
    # 2. 将用户消息加入对话历史（与原代码一致）
    conversation_history.append({"role": "user", "content": query})
    
    # 3. 构建请求消息（与原代码一致）
    messages = _build_messages_with_retention(conversation_history, retain_rounds)
    
    # 4. 定义流式生成器函数（核心新增）
    def stream_response():
        try:
            # 调用Kimi API并启用流式（stream=True）
            completion = client.chat.completions.create(
                model="moonshot-v1-8k",
                messages=messages,
                temperature=0.3,
                stream=True  # 启用流式输出
            )
            
            # 逐块处理流式响应
            full_response = ""  # 用于拼接完整响应，存入历史
            for chunk in completion:
                # 提取当前块的内容（忽略空块）
                chunk_content = chunk.choices[0].delta.content
                if chunk_content:
                    full_response += chunk_content
                    # 以SSE（Server-Sent Events）格式返回当前块（兼容前端解析）
                    yield f"data: {chunk_content}\n\n"
            
            # 流式结束后，将完整响应存入对话历史
            conversation_history.append({"role": "assistant", "content": full_response})
            
        except Exception as e:
            # 异常时返回错误信息（流式格式）
            yield f"data: Error: {str(e)}\n\n"
    
    # 5. 返回流式响应（设置MIME类型为text/event-stream）
    return Response(
        stream_response(),
        mimetype='text/event-stream',
        headers={
            "Cache-Control": "no-cache",  # 禁止缓存，确保实时性
            "X-Accel-Buffering": "no"     # 禁用反向代理缓冲（如Nginx）
        }
    )

if __name__ == '__main__':
    # 生产环境需关闭debug，避免内存泄漏
    app.run(host='0.0.0.0', port=8088, debug=False)